import json

prompt = '''{
    "system": "{}", 
    "input": {}",
    "output": {}
}'''


with open("./data/predata.json", "w") as f:
    with open("./data/alpaca_data_cleaned.json", "r") as f1:
        for line in f:
            example = json.loads(line)
            instruction = example["system"]
            input_text = example["conversation"][0]['human']
            output = example["conversation"][0]['assistant']
            text = prompt.format(instruction, input_text, output)
            f1.write(json.dumps(text)+ "\n")
                
            
with open("train.jsonl", "r") as f:
    with open("alpaca_data_cleaned.jsonl", "w") as f1:
        for line in f.readlines():
            example = json.loads(line)
            instruction = example["system"]
            input_text = example["conversation"][0]["human"]
            output = example["conversation"][0]["assistant"]
            data = {"messages": [{"role": "system", "content": instruction}, {"role": "user", "content": input_text}, {"role": "assistant", "content": output}]}
            print(data)
            f1.write(json.dumps(prompt) + '\n')
            
from datasets import load_dataset
from trl import SFTTrainer
import transformers

dataset = load_dataset("tatsu-lab/alpaca", split="train")

model = transformers.AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
tokenizer = transformers.AutoTokenizer.from_pretrained("facebook/opt-350m")

def formatting_prompts_func(examples):
    output_text = []
    for i in range(len(examples["instruction"])):
        instruction = examples["instruction"][i]
        input_text = examples["input"][i]
        response = examples["output"][i]

        if len(input_text) >= 2:
            text = f'''Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
            
            ### Instruction:
            {instruction}
            
            ### Input:
            {input_text}
            
            ### Response:
            {response}
            '''
        else:
            text = f'''Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
            
            ### Instruction:
            {instruction}
            
            ### Response:
            {response}
            '''
        output_text.append(text)
        

    return output_text

trainer = SFTTrainer(
    model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    formatting_func=formatting_prompts_func,
    max_seq_length=256,
    packing=False,
)